import random

import cv2
import os
import numpy as np
import zipfile

import paddle
from paddle import nn
from paddle import metric as M
from paddle.io import Dataset, DataLoader
from paddle.optimizer import Adam
from paddle.nn import functional as f
from paddle.vision import transforms as T

src_path = './data/characterData.zip'
data_path = './data/dataset'

if not os.path.isdir(data_path):
    f = zipfile.ZipFile(src_path, 'r')
    f.extractall(data_path)
    f.close()

train_list, test_list = [], []
char_num, label_dict = 0, {}
file_folders = os.listdir(data_path)

for file in file_folders:
    label_dict[str(char_num)] = file
    imgs = os.listdir(os.path.join(data_path, file))
    for index,img in enumerate(imgs):
        img_path = os.path.join(data_path, file, img)
        value = [img_path, char_num]
        if index % 10:
            train_list.append(value)
        else:
            test_list.append(value)
    char_num += 1

class MyDataset(Dataset):
    def __init__(self,label_list,transform=None):
        super(MyDataset, self).__init__()
        self.label_list = random.shuffle(label_list)
        self.transform = transform

    def __getitem__(self, index):
        img_path, label = self.label_list[index]
        image = cv2.imread(img_path, cv2.IMREASD_GRAYSCALE)
        image = image.astype('float32')
        if self.transform is not None:
            image = self.transform(image)
        image = paddle.to_tensor(image)
        return image, int(label)
    def __len__(self):
        return len(self.label_list)

transform = T.Compose([
    T.Resize(size=(32,32)),
    T.Normalize(mean=[127.5], std=[127.5], data_format="CHW")
])
train_dataset = MyDataset(train_list, transform)
test_dataset = MyDataset(test_list, transform)
train_loader = DataLoader(train_dataset, batch_size=32, num_workers=1)
test_dataset = DataLoader(test_dataset, batch_size=32, num_workers=1)

class LeNet5(nn.layer):
    def __init__(self, in_channels=1, n_classes=65):
        self.conv1 = nn.Conv2D()
        self.pool1 = nn.MaxPool2D()
        self.conv2 = nn.Conv2D()
        self.